# src/algorithms/calibration.py
import cv2
import numpy as np
from typing import List, Tuple
from src.data.structures import BinocularImagePair, CalibrationData


class ZhangCalibration:
    """
    张正友标定法的立体相机标定工具类

    支持棋盘格角点检测、单目标定和双目标定, 输出标定参数
    """

    def __init__(self, board_size: Tuple[int, int] = (9, 6), square_size: float = 25.0):
        """
        初始化标定器

        参数:
            board_size (Tuple[int, int]): 棋盘格的网格数, 如(9,6)
            square_size (float): 棋盘格实际边长, 单位mm
        """
        self.board_size = board_size  # 棋盘格的网格数
        self.square_size = square_size  # 棋盘格真实边长 (mm)
        self.criteria = (
            cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
            30,
            0.001,
        )  # 角点亚像素搜索迭代终止条件: EPS+MAX_ITER, maxIter=30, ε=1e-3

    def binocularCalibration(
        self, image_pairs: List[BinocularImagePair], visualize: bool = False
    ) -> CalibrationData:
        """
        执行双目相机标定流程

        参数:
            image_pairs (List[StereoImagePair]): 左右相机采集的棋盘格图像对列表
        返回:
            CalibrationData: 标定结果数据结构, 包含内参、外参、重投影误差等
        异常:
            ValueError: 有效标定图像对数量不足时抛出
        """
        # 准备标定数据
        object_points = []
        left_image_points = []
        right_image_points = []

        # 1. 生成一个平面棋盘在世界坐标系下的 3D 点集
        corner_size = (self.board_size[0] - 1, self.board_size[1] - 1)  # 内角点数
        objp = np.zeros((corner_size[0] * corner_size[1], 3), np.float32)  # 初值为 0
        # np.mgrid 生成一个网格坐标: shape=(2,8,5), 取 .T 后变 (5,8,2), reshape(-1,2) = (40,2)
        # 将前两列设置为像素坐标, 第3列维持 0 表示平面
        objp[:, :2] = np.mgrid[0 : corner_size[0], 0 : corner_size[1]].T.reshape(-1, 2)
        # 扩放到真实尺寸 (mm)
        objp *= self.square_size

        print(f"Processing {len(image_pairs)} image pairs...")
        valid_pairs = 0

        for i, pair in enumerate(image_pairs):
            # 确保输入为灰度图像
            left_gray = pair.left_image
            right_gray = pair.right_image
            if left_gray.ndim == 3:
                left_gray = cv2.cvtColor(left_gray, cv2.COLOR_BGR2GRAY)
            if right_gray.ndim == 3:
                right_gray = cv2.cvtColor(right_gray, cv2.COLOR_BGR2GRAY)

            # 粗检测棋盘角点
            ret_left, corners_left = cv2.findChessboardCorners(
                left_gray, corner_size, None
            )
            ret_right, corners_right = cv2.findChessboardCorners(
                right_gray, corner_size, None
            )

            if ret_left and ret_right:
                # 亚像素精化: 在 11×11 窗口内, 用 criteria 终止
                corners_left = cv2.cornerSubPix(
                    left_gray, corners_left, (11, 11), (-1, -1), self.criteria
                )
                corners_right = cv2.cornerSubPix(
                    right_gray, corners_right, (11, 11), (-1, -1), self.criteria
                )

                # 保存有效点集
                object_points.append(objp.astype(np.float32))
                left_image_points.append(corners_left.astype(np.float32))
                right_image_points.append(corners_right.astype(np.float32))
                valid_pairs += 1
                print(f"✅ Pair {i+1}: corners detected")

                # 绘制并显示角点
                if visualize:
                    img = cv2.drawChessboardCorners(
                        pair.left_image,
                        (corner_size[0], corner_size[1]),
                        corners_left,
                        ret_left,
                    )
                    img = self._drawCornerDirection(img, corners_left)
                    cv2.imshow("Left Corners", img)

                    img = cv2.drawChessboardCorners(
                        pair.right_image,
                        (corner_size[0], corner_size[1]),
                        corners_right,
                        ret_right,
                    )
                    img = self._drawCornerDirection(img, corners_right)
                    cv2.imshow("Right Corners", img)
                    cv2.waitKey(500)  # 显示图像 0.5 秒
            else:
                print(f"❌ Pair {i+1}: corners not found")

        if valid_pairs < 10:
            raise ValueError(
                f"Insufficient valid calibration images: {valid_pairs}/10 required"
            )

        print(f"Using {valid_pairs} valid pairs for calibration...")

        # 2. 单目标定: 左右各自标定
        shape = image_pairs[0].left_image.shape
        if len(shape) == 2:
            h, w = shape
        else:
            h, w = shape[:2]
        image_size = (w, h)  # 注意 OpenCV 要求 (width, height)

        print("Calibrating left camera...")
        ret_left, K_left, D_left, rvecs_left, tvecs_left = cv2.calibrateCamera(
            object_points, left_image_points, image_size, np.eye(3), np.zeros((5, 1))
        )
        # K_left: 3×3 内参数矩阵; D_left: (5,1) 畸变系数 [k1,k2,p1,p2,k3]
        print("Calibrating right camera...")
        ret_right, K_right, D_right, rvecs_right, tvecs_right = cv2.calibrateCamera(
            object_points, right_image_points, image_size, np.eye(3), np.zeros((5, 1))
        )

        # 3. 双目标定: 使用单目结果作为初值
        print("Performing stereo calibration...")
        # flags=cv2.CALIB_USE_INTRINSIC_GUESS: 保持并优化传入的 K 与 D, 而不是重新初始化
        # 输出 R(3×3 旋转矩阵); T(3×1 相机中心平移); E(本征矩阵); F(基础矩阵)
        ret, K_left, D_left, K_right, D_right, R, T, E, F = cv2.stereoCalibrate(
            object_points,
            left_image_points,
            right_image_points,
            K_left,
            D_left,
            K_right,
            D_right,
            image_size,
            criteria=self.criteria,
            flags=cv2.CALIB_USE_INTRINSIC_GUESS,
        )

        print(f"Stereo calibration completed. RMS error: {ret:.4f}")

        return CalibrationData(
            camera_matrix_left=K_left,
            camera_matrix_right=K_right,
            dist_coeffs_left=D_left,
            dist_coeffs_right=D_right,
            rotation_matrix=R,
            translation_vector=T,
            essential_matrix=E,
            fundamental_matrix=F,
            reprojection_error=ret,
            image_size=image_size,
        )

    def _drawCornerDirection(
        self,
        img: np.ndarray,
        corners: np.ndarray,
        color_chain=(0, 255, 255),
    ) -> np.ndarray:
        vis = img.copy()

        N = corners.shape[0]
        # 起点、终点坐标
        pt_start = tuple(corners[0].ravel().astype(int))
        pt_end = tuple(corners[-1].ravel().astype(int))

        # 1. 起点箭头（0→1）
        if N > 1:
            pt_next = tuple(corners[1].ravel().astype(int))
            cv2.arrowedLine(vis, pt_start, pt_next, color_chain, 2, tipLength=0.2)

        # 2. 终点箭头（-2→-1）
        if N > 1:
            pt_prev = tuple(corners[-2].ravel().astype(int))
            cv2.arrowedLine(vis, pt_prev, pt_end, color_chain, 2, tipLength=0.2)

        # 3. 起点标记
        cv2.circle(vis, pt_start, 6, (0, 255, 0), -1)
        cv2.putText(
            vis,
            "begin",
            (pt_start[0] - 8, pt_start[1] - 8),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 255, 0),
            2,
        )

        # 4. 终点标记
        cv2.circle(vis, pt_end, 6, (0, 0, 255), -1)
        cv2.putText(
            vis,
            "end",
            (pt_end[0] + 8, pt_end[1] + 8),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 0, 255),
            2,
        )

        return vis
